# -*- coding: UTF-8 -*-

from __future__ import absolute_import, division, print_function, unicode_literals

import datetime  # For datetime objects
import os.path  # To manage paths
import sys  # To find out the script name (in argv[0])

# Import the backtrader platform
import backtrader as bt

import MySQLDataFeed as mdf
import SimpleMySqlClass as smc
import talib as talib
import numpy as np
import akshare as ak
import pandas as pd
import json

"""
双移动平均线​​（50日/200日）判断趋势方向
​​ADX指标​​（14日周期）判断趋势强度
当ADX>25且短期均线在长期均线上方时为牛市
当ADX>25且短期均线在长期均线下方时为熊市
ADX<25时判定为震荡市

如果ADX > 25，且MA50 > MA200，则为牛市。
如果ADX > 25，且MA50 < MA200，则为熊市。
否则，判断为震荡市。

"""


class MarketPhaseStrategy(bt.Strategy):
    params = (
        ("ma_short", 50),
        ("ma_long", 200),
        ("adx_period", 14),
        ("adx_threshold", 25),
    )

    def __init__(self):
        # 初始化技术指标
        self.ma_short = bt.indicators.SMA(self.data.close, period=self.p.ma_short)
        self.ma_long = bt.indicators.SMA(self.data.close, period=self.p.ma_long)
        self.adx = bt.indicators.AverageDirectionalMovementIndex(
            self.data, period=self.p.adx_period
        )
        # 增加布林带和波动率判断
        self.boll = bt.indicators.BollingerBands(self.data.close)
        self.atr = bt.indicators.ATR(self.data)

        # 存储市场阶段和日期
        self.market_phases = []

    def next(self):
        # 确保有足够的数据长度
        if len(self) < max(self.p.ma_long, self.p.adx_period):
            return

        current_date = self.data.datetime.date()
        phase = "---"
        # 判断市场阶段
        if self.adx[0] > self.p.adx_threshold:
            if self.ma_short[0] > self.ma_long[0]:
                phase = "牛市"
            else:
                phase = "熊市"
        else:
            phase = "震荡市"
            # 增加震荡市判断条件
            # if (self.boll.lines.top[0] - self.boll.lines.bot[0]) < (
            #     self.data.close[0] * 0.05
            # ) and self.atr[0] < self.data.close[0] * 0.01:
            #     phase = "震荡市"

        self.market_phases.append((current_date, phase))
        print(f"日期: {current_date}, 市场阶段: {phase}")


# 回测设置
def run_backtest():
    cerebro = bt.Cerebro()

    code = "510300"
    name = "沪深300ETF"

    tableName = "t_trade_data_" + code
    sql = f"select `date`, `open`, `high`, `low`, `close`, `volume` from {tableName} where date>'2022-01-01' limit 3000"

    data = mdf.MySQLDataFeed(sql)
    print("data = ", data)
    cerebro.adddata(data)
    cerebro.addstrategy(MarketPhaseStrategy)
    cerebro.run()

    return cerebro


if __name__ == "__main__":
    cerebro = run_backtest()

    # 提取市场阶段数据
    strategy = cerebro.runstrats[0][0]
    phase_df = pd.DataFrame(strategy.market_phases, columns=["date", "phase"])

    # 输出最后10天的市场阶段
    print("\n最近市场阶段分析:")
    print(phase_df.tail(10))
